package com.yzh.zhipu.session.impl;

import com.yzh.zhipu.ChatGlmApi;
import com.yzh.zhipu.executor.Executor;
import com.yzh.zhipu.model.ChatRequest;
import com.yzh.zhipu.model.ChatResponse;
import com.yzh.zhipu.model.ImageRequest;
import com.yzh.zhipu.model.ImageResponse;
import com.yzh.zhipu.model.vo.Model;
import com.yzh.zhipu.session.Configuration;
import com.yzh.zhipu.session.ChatGlmSession;
import lombok.extern.slf4j.Slf4j;
import okhttp3.sse.EventSource;
import okhttp3.sse.EventSourceListener;

import java.io.IOException;
import java.util.HashMap;
import java.util.concurrent.CompletableFuture;

@Slf4j
public class ChatGlmSessionImpl implements ChatGlmSession {
    private Configuration configuration;
    private HashMap<Model, Executor> executors;
    private ChatGlmApi chatGlmApi;

    public ChatGlmSessionImpl(Configuration configuration, HashMap<Model, Executor> executors) {
        this.configuration = configuration;
        this.executors = executors;
        this.chatGlmApi = configuration.getChatGlmApi();
    }

    @Override
    public EventSource completions(ChatRequest chatRequest, EventSourceListener eventSourceListener) {
        Executor executor = executors.get(chatRequest.getModel());
        if (executor == null) throw new RuntimeException(chatRequest.getModel() + " 模型执行器尚未实现！");
        return executor.completions(chatRequest, eventSourceListener);
    }

    @Override
    public ChatResponse completionsSync(ChatRequest chatRequest) throws IOException {
        Executor executor = executors.get(chatRequest.getModel());
        if (executor == null) throw new RuntimeException(chatRequest.getModel() + " 模型执行器尚未实现！");
        return executor.completionsSync(chatRequest);
    }

    @Override
    public CompletableFuture<java.lang.String> completions(ChatRequest chatRequest) {
        Executor executor = executors.get(chatRequest.getModel());
        if (executor == null) throw new RuntimeException(chatRequest.getModel() + " 模型执行器尚未实现！");
        return executor.completions(chatRequest);
    }

    @Override
    public ImageResponse genImage(ImageRequest imageRequest) throws IOException {
        Executor executor = executors.get(imageRequest.getModel());
        if (executor == null) throw new RuntimeException(imageRequest.getModel() + " 模型执行器尚未实现！");
        return executor.genImage(imageRequest);
    }

}
